from typing import Sized
from matplotlib.pyplot import close, plot
import pandas as pd
import backtrader as bt
from backtrader_plotting import Bokeh
from backtrader_plotting.schemes import Tradimo
from my_plot import *
from py_mysql import *
import numpy as np

class PandasData_more(bt.feeds.PandasData):
    pass

# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )
    ema_periodArr = [12,144,169,576,676]

    # 初始化部分指标

    def __init__(self):
        self.dataclose = self.datas[0].close
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None
        # 加入均线指标
        self.ema12 = bt.indicators.ExponentialMovingAverage(self.datas[0].close, period=self.ema_periodArr[0])
        self.ema144 = bt.indicators.ExponentialMovingAverage(self.datas[0].close, period=self.ema_periodArr[1])
        self.ema169 = bt.indicators.ExponentialMovingAverage(self.datas[0].close, period=self.ema_periodArr[2])
        self.ema576 = bt.indicators.ExponentialMovingAverage(self.datas[0].close, period=self.ema_periodArr[3])
        self.ema676 = bt.indicators.ExponentialMovingAverage(self.datas[0].close, period=self.ema_periodArr[4])

        self.maxClose = bt.indicators.Highest(self.datas[0].close, period=60)
        self.minClose = bt.indicators.Lowest(self.datas[0].close, period=60)

        # 止盈止损点
        self.topPrice = 0
        self.botPrice = 0

        # 止盈止损比例
        self.stopProfit = 1.382
        self.stopLoss = 0.618
        # 下单数量
        self.volSize = 10


    # 策略核心代码
    def next(self):
        # 当前日期
        self.time = str(bt.num2date(self.lines.datetime[0]))
        self.positionSize = self.position.size

        close = self.dataclose[0]

        # 当前空仓
        if not self.position:
            if (self.dataclose[0] > self.ema12[0]) and (self.ema144[0] > self.ema169[0]) and (self.ema576[0] > self.ema676[0]) and (self.ema12[0] > self.ema144[0]) and (self.ema169[0] > self.ema576[0]):
                if (self.dataclose[-1] < self.ema12[-1]):
                    self.buy(size = self.volSize)
                    bad_size = close - self.minClose[0]
                    self.topPrice = close + bad_size*self.stopProfit
                    self.botPrice = close - bad_size*self.stopLoss
            elif (self.dataclose[0] < self.ema12[0]) and (self.ema144[0] < self.ema169[0]) and (self.ema576[0] < self.ema676[0]) and (self.ema12[0] < self.ema144[0]) and (self.ema169[0] < self.ema576[0]):
                # 开仓做空
                if (self.dataclose[-1] > self.ema12[-1]):
                    self.sell(size = self.volSize)
                    bad_size = self.maxClose[0] - close
                    self.topPrice = close + bad_size*self.stopLoss
                    self.botPrice = close - bad_size*self.stopProfit
        # 有持仓
        else:
            if close > self.topPrice or close < self.botPrice:
                    print('触碰止盈止损,进行平仓')
                    self.close()

        

    # 日志输出
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                print('日期：{}---买入价格:{}---买入手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                print('日期：{}---卖出价格:{}---卖出手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return

    # 策略结束时，多用于参数调优
    def stop(self):
        print(self.positionSize)
        self.log('(MA均线： %2d日) 期末总资金 %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=False)


def main(startcash=1000000, com=0.0003, qts=100):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    # df = selectSotck(adjustflag, code, start, end)
    # df = writeExcl()  
    # df.index = pd.to_datetime(df.date)
    # df = df[['open', 'high', 'low', 'close', 'volume']]
    # # 将数据加载至回测系统
    # data = bt.feeds.PandasData(dataname=df)

    # cerebro.adddata(data,name='RBL8')
    # 获取数据
    query_db = Mysql_search()
    df = query_db.get_one(['RB'],'2020-06-01','2020-08-31')
   
    for item in df:
        df2 = df[item]
        df2.index = pd.to_datetime( df2.index, utc= True)
        print(df2)
        data = PandasData_more(
            dataname = df2,
            timeframe=bt.TimeFrame.Minutes)
        cerebro.adddata(data)
    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(False)

    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')
    result = cerebro.run()
    print('期末总资金: %.2f' % cerebro.broker.getvalue())

    # custom_plot(result)
    cerebro.plot()
    
    b = Bokeh(style='bar', tabs='multi', scheme=Tradimo())
    # cerebro.plot(b)

def writeExcl():
    df = pd.read_excel(r'D:\通达信\T0002\export2\IL8.xls')
    df['code'] = 'IL8'
    return df

main()

